//ITMD536 Software Testing and Maintenance.
//Final project <group assignment>
//Author: Wanglei ID:A20479012
#include <iostream>
#include <string>
#include <vector>
#include <algorithm>
#include <map>
#include "Face.hpp"

const float threshold[3] = { 0.8f, 0.8f, 0.6f };
int minsize = 40;
const float pre_facetor = 0.709f;
const float nms_threshold[3] = {0.5f, 0.7f, 0.7f};
const float mean_vals[3] = {127.5, 127.5, 127.5};
const float norm_vals[3] = {0.0078125, 0.0078125, 0.0078125};
const int MIN_DET_SIZE = 12;
int img_w, img_h;
ncnn::Mat img;
ncnn::Net pNet, rNet, oNet;//mtcnn algorithm p,r,o net for face detect.
ncnn::Net mNet;//for face recognition.
ncnn::Net lmNet;//for face landmark.
std::vector<Bbox> firstPreviousBbox_, secondPreviousBbox_, thirdPrevioussBbox_;
std::vector<Bbox> firstBbox_, secondBbox_,thirdBbox_;

//default contructor.
Face::Face(){
    //pnet, rnet, onet model load ...
	pNet.load_param("/usr/share/model/det1.param");
	pNet.load_model("/usr/share/model/det1.bin");
	rNet.load_param("/usr/share/model/det2.param");
	rNet.load_model("/usr/share/model/det2.bin");
	oNet.load_param("/usr/share/model/det3.param");
	oNet.load_model("/usr/share/model/det3.bin");

    //load mfn models.
	mNet.load_param("/usr/share/model/mfn-4Mb-v2.0-fp32.param");
	mNet.load_model("/usr/share/model/mfn-4Mb-v2.0-fp32.bin");
	
    //load landmark models.
    lmNet.load_param("/usr/share/model/landmark.param");
	lmNet.load_model("/usr/share/model/landmark.bin");
	std::cout << "face contructor: load models success, done." << std::endl;
}

bool cmpScore(Bbox lsh, Bbox rsh) 
{
	if (lsh.score < rsh.score)
		return true;
	else
		return false;
}

static void generateBbox(ncnn::Mat score, ncnn::Mat location, std::vector<Bbox>& boundingBox_, float scale)
{
	const int stride = 2;
	const int cellsize = 12;
	//score p
	float *p = score.channel(1);//score.data + score.cstep;
	//float *plocal = location.data;
	Bbox bbox;
	float inv_scale = 1.0f/scale;

	for(int row=0;row<score.h;row++)
	{
		for(int col=0;col<score.w;col++)
		{
			if(*p>threshold[0])
			{
				bbox.score = *p;
				bbox.x1 = round((stride*col+1)*inv_scale);
				bbox.y1 = round((stride*row+1)*inv_scale);
				bbox.x2 = round((stride*col+1+cellsize)*inv_scale);
				bbox.y2 = round((stride*row+1+cellsize)*inv_scale);
				bbox.area = (bbox.x2 - bbox.x1) * (bbox.y2 - bbox.y1);
				const int index = row * score.w + col;

				for(int channel=0;channel<4;channel++)
					bbox.regreCoord[channel]=location.channel(channel)[index];
				boundingBox_.push_back(bbox);
			}

			p++; //plocal++;
		}
	}
	
	return;
}

static void nms(std::vector<Bbox> &boundingBox_, const float overlap_threshold, std::string modelname)
{
	if(boundingBox_.empty()) return;
	sort(boundingBox_.begin(), boundingBox_.end(), cmpScore);
	float IOU = 0;
	float maxX = 0;
	float maxY = 0;
	float minX = 0;
	float minY = 0;
	std::vector<int> vPick;
	int nPick = 0;
	std::multimap<float, int> vScores;
	const int num_boxes = boundingBox_.size();
	vPick.resize(num_boxes);
	
	for (int i = 0; i < num_boxes; ++i)
		vScores.insert(std::pair<float, int>(boundingBox_[i].score, i));
	while(vScores.size() > 0)
	{
		int last = vScores.rbegin()->second;
		vPick[nPick] = last;
		nPick += 1;
		for (std::multimap<float, int>::iterator it = vScores.begin(); it != vScores.end();)
		{
			int it_idx = it->second;
			maxX = std::max(boundingBox_.at(it_idx).x1, boundingBox_.at(last).x1);
			maxY = std::max(boundingBox_.at(it_idx).y1, boundingBox_.at(last).y1);
			minX = std::min(boundingBox_.at(it_idx).x2, boundingBox_.at(last).x2);
			minY = std::min(boundingBox_.at(it_idx).y2, boundingBox_.at(last).y2);
			//maxX1 and maxY1 reuse 
			maxX = ((minX-maxX+1)>0)? (minX-maxX+1) : 0;
			maxY = ((minY-maxY+1)>0)? (minY-maxY+1) : 0;
			//IOU reuse for the area of two bbox
			IOU = maxX * maxY;
			if(!modelname.compare("Union"))
			IOU = IOU/(boundingBox_.at(it_idx).area + boundingBox_.at(last).area - IOU);
			else if(!modelname.compare("Min"))
				IOU = IOU/((boundingBox_.at(it_idx).area < boundingBox_.at(last).area)? boundingBox_.at(it_idx).area : boundingBox_.at(last).area);
			if(IOU > overlap_threshold)
				it = vScores.erase(it);
			else
				it++;
		}
	}

	vPick.resize(nPick);
	std::vector<Bbox> tmp_;
	tmp_.resize(nPick);
	
	for(int i = 0; i < nPick; i++)
		tmp_[i] = boundingBox_[vPick[i]];
	boundingBox_ = tmp_;
	return;
}

static void refine(std::vector<Bbox> &vecBbox, const int &height, const int &width, bool square)
{
	if(vecBbox.empty())
	{
		std::cout<<"Bbox is empty!!"<<std::endl;
		return;
	}

	float bbw=0, bbh=0, maxSide=0;
	float h = 0, w = 0;
	float x1=0, y1=0, x2=0, y2=0;
	for(std::vector<Bbox>::iterator it=vecBbox.begin(); it!=vecBbox.end();it++)
	{
		bbw = (*it).x2 - (*it).x1 + 1;
		bbh = (*it).y2 - (*it).y1 + 1;
		x1 = (*it).x1 + (*it).regreCoord[0]*bbw;
		y1 = (*it).y1 + (*it).regreCoord[1]*bbh;
		x2 = (*it).x2 + (*it).regreCoord[2]*bbw;
		y2 = (*it).y2 + (*it).regreCoord[3]*bbh;

		if(square)
		{
			w = x2 - x1 + 1;
			h = y2 - y1 + 1;
			maxSide = (h>w)?h:w;
			x1 = x1 + w*0.5 - maxSide*0.5;
			y1 = y1 + h*0.5 - maxSide*0.5;
			(*it).x2 = round(x1 + maxSide - 1);
			(*it).y2 = round(y1 + maxSide - 1);
			(*it).x1 = round(x1);
			(*it).y1 = round(y1);
		}

		//boundary check
		if((*it).x1<0)(*it).x1=0;
		if((*it).y1<0)(*it).y1=0;
		if((*it).x2>width)(*it).x2 = width - 1;
		if((*it).y2>height)(*it).y2 = height - 1;

		it->area = (it->x2 - it->x1)*(it->y2 - it->y1);
	}

	return;
}

static void rpNet(void)
{
	firstBbox_.clear();
	float minl = img_w < img_h? img_w: img_h;
	float m = (float)MIN_DET_SIZE/minsize;
	minl *= m;
	float factor = pre_facetor;
	std::vector<float> scales_;
	while(minl>MIN_DET_SIZE)
	{
		scales_.push_back(m);
		minl *= factor;
		m = m*factor;
	}
	
	for (size_t i = 0; i < scales_.size(); i++) 
	{
		int hs = (int)ceil(img_h*scales_[i]);
		int ws = (int)ceil(img_w*scales_[i]);
		ncnn::Mat in;
		resize_bilinear(img, in, ws, hs);
		ncnn::Extractor ex = pNet.create_extractor();
		//ex.set_num_threads(2);
		ex.set_light_mode(true);
		ex.input("data", in);
		ncnn::Mat score_, location_;
		ex.extract("prob1", score_);
		ex.extract("conv4-2", location_);
		std::vector<Bbox> boundingBox_;
		generateBbox(score_, location_, boundingBox_, scales_[i]);
		nms(boundingBox_, nms_threshold[0], "Union");
		firstBbox_.insert(firstBbox_.end(), boundingBox_.begin(), boundingBox_.end());
		boundingBox_.clear();
	}
	
	return;
}

static void rrNet(void)
{
	secondBbox_.clear();
	//int count = 0;

	for(std::vector<Bbox>::iterator it=firstBbox_.begin(); it!=firstBbox_.end();it++)
	{
		ncnn::Mat tempIm;
		copy_cut_border(img, tempIm, (*it).y1, img_h-(*it).y2, (*it).x1, img_w-(*it).x2);
		ncnn::Mat in;
		resize_bilinear(tempIm, in, 24, 24);
		ncnn::Extractor ex = rNet.create_extractor();
		//ex.set_num_threads(2);
		ex.set_light_mode(true);
		ex.input("data", in);
		ncnn::Mat score, bbox;
		ex.extract("prob1", score);
		ex.extract("conv5-2", bbox);

		if ((float)score[1] > threshold[1]) 
		{
			for (int channel = 0; channel<4; channel++) 
				it->regreCoord[channel] = (float)bbox[channel];//*(bbox.data+channel*bbox.cstep);
			it->area = (it->x2 - it->x1)*(it->y2 - it->y1);
			it->score = score.channel(1)[0];//*(score.data+score.cstep);
			secondBbox_.push_back(*it);
		}
	}

	return;
}

static void roNet(void)
{
	thirdBbox_.clear();
	for(std::vector<Bbox>::iterator it=secondBbox_.begin(); it!=secondBbox_.end();it++)
	{
		ncnn::Mat tempIm;
		copy_cut_border(img, tempIm, (*it).y1, img_h-(*it).y2, (*it).x1, img_w-(*it).x2);
		ncnn::Mat in;
		resize_bilinear(tempIm, in, 48, 48);
		ncnn::Extractor ex = oNet.create_extractor();
		//ex.set_num_threads(2);
		ex.set_light_mode(true);
		ex.input("data", in);
		ncnn::Mat score, bbox, keyPoint;
		ex.extract("prob1", score);
		ex.extract("conv6-2", bbox);
		ex.extract("conv6-3", keyPoint);
		if ((float)score[1] > threshold[2]) 
		{
			for (int channel = 0; channel < 4; channel++) 
				it->regreCoord[channel] = (float)bbox[channel];
			
			it->area = (it->x2 - it->x1) * (it->y2 - it->y1);
			it->score = score.channel(1)[0];
			for(int num = 0; num<5; num++) 
			{
				//(it->ppoint)[num] = it->x1 + (it->x2 - it->x1) * keyPoint[num];
				it->landmark.x[num] = it->x1 + (it->x2 - it->x1) * keyPoint[num];
				//(it->ppoint)[num + 5] = it->y1 + (it->y2 - it->y1) * keyPoint[num + 5];
				it->landmark.y[num] = it->y1 + (it->y2 - it->y1) * keyPoint[num + 5];
			}

			thirdBbox_.push_back(*it);
		}
	}

	return;
}

int Face::faceDetect(cv::Mat& frame, std::vector<Bbox>& rBbox){
    //cv Mat to ncnn Mat.
	ncnn::Mat nImg = ncnn::Mat::from_pixels(frame.data, ncnn::Mat::PIXEL_BGR2RGB, frame.cols, frame.rows);
	img = nImg;
	img.substract_mean_normalize(mean_vals, norm_vals);
	img_w = img.w;
	img_h = img.h;
	
	//1. first stage, run the pnet ...
	rpNet();
	if(firstBbox_.size() < 1) 
		return facePnetError;
	nms(firstBbox_, nms_threshold[0], "Union");
	refine(firstBbox_, img_h, img_w, true);
	//printf("firstBbox_.size()=%d\n", firstBbox_.size());

	//2. second stage, run the rnet ...
	rrNet();
	if(secondBbox_.size() < 1) 
		return faceRnetError;
	nms(secondBbox_, nms_threshold[1], "Union");
	refine(secondBbox_, img_h, img_w, true);
	//printf("secondBbox_.size()=%d\n", secondBbox_.size());
	
	//3. third stage, run the onet ...
	roNet();
	if(thirdBbox_.size() < 1)
		return faceOnetError;
	refine(thirdBbox_, img_h, img_w, true);
	nms(thirdBbox_, nms_threshold[2], "Min");
	rBbox = thirdBbox_;
	return faceNormal;
}

